
from math import sqrt

from PyQt4 import QtCore, QtGui
from PyQt4.QtCore import Qt
from PyQt4.QtOpenGL import *
from OpenGL.GLU import *
from OpenGL.GL import *
import numpy

from amt.cg.core.object3dlist import Object3DList
from amt.cg.core.camera import Camera
from amt.cg.core.arcball import Arcball
from amt.cg.core.util import multiply_vector_by_matrix
from amt.cg.geometry.sphere import Sphere
from amt.cg.geometry.cube import Cube
from amt.cg.ui.mouse import Mouse

class GlWidget (QGLWidget):
    '''A simple OpenGL pane'''
    
    def __init__ (self, parent=None):
        '''Constructor'''
        super (GlWidget, self).__init__(parent)
        self.current_width = 1
        self.current_height = 1
        self._camera = Camera(self.current_width, self.current_height)
        self._init_objects()
        self._arcball_sphere = Arcball("scene")
        self._mouse = Mouse()
        self._selected = None
        self._object_size = 2.0
        
    def initializeGL(self):
        '''
        Called before the screen be painted for the first time.
        '''
        glClearColor(0.75, 0.85, 0.85, 1.0)
        glEnable(GL_DEPTH_TEST)
        self._lights()
        self._camera.adjust_lens()
            
    def paintGL(self, mode=GL_RENDER):
        '''
        Sets the camera and draws all objects at the screen.
        '''
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        self._camera.look_at()
        self._object_list.draw_all(mode)
        
    def resizeGL (self, width, height):
        '''
        The resize callback.
        '''
        self.current_width, self.current_height = width, height
        self._camera.update_lens_aspect(self.current_width, self.current_height)
        
        # Calculates the arcball center and radius
        self._arcball_sphere.center = [width/2, height/2, 0]
        self._arcball_sphere.radius = sqrt((width/2)*(width/2) + (height/2)*(height/2))
        
        glViewport(0, 0, width, height)
        self.updateGL()

    # Begin of Qt event handlers

    def keyPressEvent(self, event):
        '''
        Handles the keyboard press events.
        Draws an object under the mouse cursor.
        '''
        # Deletes the selected object
        if event.text() == "x":
            self.remove_object()
            return
            
        # Centralizes the scene
        if event.key() == Qt.Key_Home:
            self.center_scene()
            return            

        xpos, ypos = self._get_window_coords()
        zpos = 0.8
        
        coords = gluUnProject(xpos, ypos, zpos)

        # Draws a cube 
        if event.text() == "c":   
            size = float(self._object_size/10.0)
            new_object = Cube(size, coords, self._camera)
            self._object_list.append(new_object)            
        
        # Draws a sphere
        elif event.text() == "e":
            size = float(self._object_size/10.0)
            size = size / sqrt(3)
            new_object = Sphere(size, coords, self._camera)
            self._object_list.append(new_object)
        
        self.updateGL()
        
    def mouseMoveEvent(self, event):
        '''
        Handles the mouse move events.
        Left button - translation
        Right button - rotation
        Middle button - zoom
        '''
        vector = self._mouse.move(event)
        
        # Translation. 
        if self._mouse.button == Mouse.LEFT:
            
            # Translation of selected objects.           
            if self._selected is not None:
                vector = self._get_translation_vector(vector)
                
                # Softens the translation
                for i in range(3):
                    vector[i] *= self._camera.movement_multiplier()
                                    
                self._selected.translate(*vector, w=0)
            
            # Translation of the scene.    
            else:
                vector[0] *= -1
                self._camera.move(vector)
        
        # Arcball rotation.
        elif self._mouse.button == Mouse.RIGHT:
            xpos, ypos = self._get_window_coords()
            
            # Rotation of selected objects
            if self._selected is not None:
                self._selected.rotate(xpos, ypos)
                self._selected.update_arcball_points()

            # Rotation of the scene
            else:
                self._arcball_sphere.set_final_point(xpos, ypos)
                self.rotate_scene()
                
                # Updating the Arcball's points
                for i in range(3):
                    self._arcball_sphere.initial_point[i] = self._arcball_sphere.final_point[i]
        
        # Zoom
        elif self._mouse.button == Mouse.MIDDLE:
            self._camera.zoom(-vector[1])

        self.updateGL()
        
    def mousePressEvent(self, event):
        '''
        Handles the mouse click events.
        With a left click, selects an object at the screen.
        '''
        self._mouse.press(event)
        xpos, ypos = self._get_window_coords()
        
        if self._mouse.button == Mouse.LEFT:
            if self._selected is not None:
                self._selected.deselect()
                
            self._selected = self._pick(xpos, ypos)
            
            if self._selected is not None:
                self._selected.select()
            
            self.updateGL()

        # Sets the arcball's initial position 
        elif self._mouse.button == Mouse.RIGHT:
            if self._selected is not None:
                self._selected.right_button_click(xpos, ypos)
            else:
                self._arcball_sphere.set_initial_point(xpos, ypos)

    def mouseReleaseEvent(self, event):
        self._mouse.release()
        
    def enterEvent(self, event):
        self.setFocus()
        self.grabKeyboard()

    def leaveEvent(self, event):
        self.releaseKeyboard()
            
    # End of Qt event handlers

    def rotate_scene(self):
        '''
        This method does the arcball rotation to the scene.
        '''
        (rotation_vector, angle) = self._arcball_sphere.rotate()
            
        glMatrixMode(GL_MODELVIEW)
        glPushMatrix()
        glLoadIdentity()
        glRotate(angle, rotation_vector[0], rotation_vector[1], rotation_vector[2])
        glMultMatrixf(self._camera.rotation_matrix)
        self._camera.rotation_matrix = glGetDouble(GL_MODELVIEW_MATRIX)
        glPopMatrix()

    def remove_object(self):
        '''
        This method removes the selected object.
        '''
        if self._selected is not None:
            self._object_list.remove(self._selected)
            self.updateGL()
    
    def clear_scene(self):
        '''
        This methos removes all objects.
        '''
        self._object_list = Object3DList()
        self._camera.reset()
        self.updateGL()

    def tick_zoom(self, value):
        self._camera.zoom(value)
        self.updateGL()
        
    def center_scene(self):
        if len(self._object_list) > 0:
            box = self._object_list.get_bounding_box()
            self._camera.center_scene(box)
            self.updateGL()
        
    def _pick(self, x, y):
        '''
        Selects an object at (x, y) in window coordinates. 
        '''
        if len(self._object_list) > 0:  
            glSelectBuffer(len(self._object_list) * 4)
            
            glRenderMode(GL_SELECT)
            glInitNames()
            glPushName(0)
            
            glMatrixMode(GL_PROJECTION)
            glPushMatrix()
            glLoadIdentity()
            
            gluPickMatrix(x, y, 1, 1, glGetInteger(GL_VIEWPORT))
            self._camera.adjust_lens()
    
            glMatrixMode(GL_MODELVIEW)
            self.paintGL(GL_SELECT)
            
            glMatrixMode(GL_PROJECTION)
            glPopMatrix()
    
            self.swapBuffers()
            
            hits = glRenderMode(GL_RENDER)
            if len(hits) > 0:
                name = hits[0][2][0]
                return self._object_list[name]

        return None
    
    def _get_window_coords(self):
        '''
        Returns the cursor's location in window coordinates.
        The origin is at the bottom left corner of the viewport.
        '''        
        x = self.mapFromGlobal(QtGui.QCursor.pos()).x()
        y = self.mapFromGlobal(QtGui.QCursor.pos()).y()

        viewport = glGetInteger(GL_VIEWPORT)
        y = viewport[3] - y - 1

        return x, y
    
    def _get_translation_vector(self, vector):
        '''
        Returns a translation vector for the selected object
        in the correct coordinates system.
        '''
        modelview = glGetDouble(GL_MODELVIEW_MATRIX)
        inverse = numpy.linalg.inv(modelview)
        
        glPushMatrix()
        glLoadMatrixf(inverse)
        
        vector[1] *= -1
        result = multiply_vector_by_matrix(vector)
        glPopMatrix()
        
        return result
        
    def _lights(self):
        '''
        Sets up the scene's lights.
        '''
        glEnable(GL_LIGHTING)
        lightZeroPosition = [4.0, 4.0, 10.0, 1.0]
        lightZeroColor = [0.6, 0.6, 0.6, 1.0]
        glLightfv(GL_LIGHT0, GL_POSITION, lightZeroPosition)
        glLightfv(GL_LIGHT0, GL_DIFFUSE, lightZeroColor)
        
        glLightfv(GL_LIGHT0, GL_AMBIENT, [0.0, 0.5, 0.5, 1.0])
        glLightf(GL_LIGHT0, GL_CONSTANT_ATTENUATION, 0.1)
        glLightf(GL_LIGHT0, GL_LINEAR_ATTENUATION, 0.05)
        glEnable(GL_LIGHT0)

    def _init_objects(self):
        '''
        Initializes the 3D object list.
        '''          
        self._object_list = Object3DList()

    # The methods below are used by the interface to change the object color
    def set_color(self, color):
        if self._selected is not None:
            self._selected.color = color        
            self.updateGL()
        
    def set_yellow_color(self):
        self.set_color([1.0, 1.0, 0.0])           
            
    def set_orange_color(self):
        self.set_color([1.0, 0.67, 0.0])

    def set_red_color(self):
        self.set_color([1.0, 0.0, 0.0])

    def set_pink_color(self):
        self.set_color([1.0, 0.0, 1.0])

    def set_green_color(self):
        self.set_color([0.0, 0.67, 0.0])

    def set_purple_color(self):
        self.set_color([0.67, 0.0, 1.0])

    def set_gray_color(self):
        self.set_color([0.52, 0.52, 0.52])
